package com.gitee.wsl.flow.event

import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.distinctUntilChanged
import kotlinx.coroutines.flow.dropWhile
import kotlinx.coroutines.flow.map
import kotlinx.coroutines.flow.withIndex

enum class SharedFlowCollectionEvent {
    NoMoreCollections,
    FirstCollection,
    LaterCollections,
}

/**
 * Calls back with events in the following order:
 *
 * [FirstCollection] -> [NoMoreCollections] -> [LaterCollections] -> back to [NoMoreCollections] when all flows are terminated
 *
 * If a flow already exists on this [MutableSharedFlow] the [FirstCollection] event is immediately raised.
 * If no subscriptions exist
 *
 * Note that if this [MutableSharedFlow] was subscribed and unsubscribed before this method was called,
 * a new subscription will still raise the [FirstCollection] event
 */
suspend fun <T> MutableSharedFlow<T>.onCollectionEvent(
    onEvent: suspend MutableSharedFlow<T>.(SharedFlowCollectionEvent) -> Unit
): Unit = subscriptionCount
    .map { it > 0 } // simplify to zero (false) or one-or-more (true)
    .dropWhile { !it } // do not report until the first collection
    .distinctUntilChanged() // report as state not an event
    .withIndex() // add an index so we distinguish the first event
    .collect { subscriptionState ->
        onEvent(
            when {
                subscriptionState.value && subscriptionState.index == 0 -> SharedFlowCollectionEvent.FirstCollection
                subscriptionState.value && subscriptionState.index != 0 -> SharedFlowCollectionEvent.LaterCollections
                else -> SharedFlowCollectionEvent.NoMoreCollections
            },
        )
    }
